#include <pthread.h>
#include <arpa/inet.h>
#include <netinet/in.h>
#include <stdio.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <unistd.h>
#include <stdlib.h>
#include <string.h>
#include <sys/timeb.h>
#include <libssh/libssh.h>
#include <iostream>
#include <vector>
#include <sys/timeb.h>
#include <libssh/libssh.h>
#include <syslog.h>
#include <sstream>
#include <fstream>
#include <signal.h>
#include<limits.h>
#include "configfilereader.h"
#include "connection.h"
#include "attackleases.h"


#ifdef LIBSSH_VERSION_INT
#if LIBSSH_VERSION_INT < SSH_VERSION_INT (0, 5, 0)
#define ssh_string_free string_free
#define ssh_channel_new channel_new
#define ssh_channel_open_session channel_open_session
#define ssh_channel_send_eof channel_send_eof
#define ssh_channel_close channel_close
#define ssh_channel_free channel_free
#define ssh_channel_request_exec channel_request_exec
#define ssh_channel_read channel_read
#endif
#endif


#define BUFLEN 512
#define connectionTimeout 120 //in seconds
#define syncInterval 60 //in seconds

string configFile="/etc/gkmtfirewall.conf";

configfilereader *conf;
connection *c;
attackleases *al;
int newConnectionAdded=0;
int newAttackAdded=0;
vector<string> attacker;
pthread_t attackCheckerThread, attackBlockerThread, oldConnectionCheckerThread, leaseRemoverThread;
pthread_cond_t newConnectionCond  = PTHREAD_COND_INITIALIZER;
pthread_cond_t newAttackCond  = PTHREAD_COND_INITIALIZER;
pthread_mutex_t mutex1 = PTHREAD_MUTEX_INITIALIZER;
pthread_mutex_t mutex2 = PTHREAD_MUTEX_INITIALIZER;

using namespace std;

string itos(int i)
{
    std::stringstream out;
    out << i;
    return out.str();
}

void log(int level,string message,int w)
{
    if (level>conf->getLogLevel())
        return;

    if (w==0) //Error
    {
        syslog(LOG_ERR,"%s",message.c_str());
    }

    if (w==1)
    {
        syslog(LOG_NOTICE,"%s",message.c_str());
    }
}

vector<string> getFirewallAddressListIPs()
{
    vector<string> retVal;

    ssh_session my_ssh_session;
    int rc;
    my_ssh_session = ssh_new();
    if (my_ssh_session == NULL)
        return retVal;
    ssh_options_set(my_ssh_session, SSH_OPTIONS_HOST, conf->getFirewallIP().c_str());
    ssh_options_set(my_ssh_session, SSH_OPTIONS_USER, conf->getFirewallUser().c_str());
    rc = ssh_connect(my_ssh_session);
    if (rc != SSH_OK)
    {
        string msg="SSH Connection error";
        log(1,msg,0);
        return retVal;
    }
    rc = ssh_userauth_password(my_ssh_session, NULL, conf->getFirewallPass().c_str());
    if (rc != SSH_AUTH_SUCCESS)
    {
        string msg="SSH authentication error";
        log(1,msg,0);
        return retVal;
    }
    ssh_channel channel=ssh_channel_new(my_ssh_session);
    rc = ssh_channel_open_session(channel);
    if (rc != SSH_OK)
    {
        string msg="SSH channel error";
        log(1,msg,0);
        ssh_channel_free(channel);
        return retVal;
    }

    string cmd="/ip firewall address-list print where list=" + conf->getFirewallAddressList();
    rc = ssh_channel_request_exec(channel, cmd.c_str());
    if (rc != SSH_OK)
    {
        string msg="SSH command execution error (" + cmd + ")";
        log(1,msg,0);
        ssh_channel_close(channel);
        ssh_channel_free(channel);
        return retVal;
    }
    else
    {
        string msg="SSH executed : " + cmd;
        log(3,msg,1);

        char buffer[256];
        unsigned int nbytes;
        ofstream tempFile;
        tempFile.open ("tempFile.txt", ios::out);
        nbytes = ssh_channel_read(channel, buffer, sizeof(buffer)-1, 0);
        while (nbytes > 0)
        {
            buffer[nbytes]='\0';
            string ts=string(buffer);
            tempFile << buffer;
            nbytes = ssh_channel_read(channel, buffer, sizeof(buffer)-1, 0);
        }
        tempFile.close();
    }
    ssh_channel_send_eof(channel);
    ssh_channel_close(channel);
    ssh_channel_free(channel);
    ssh_disconnect(my_ssh_session);
    ssh_free(my_ssh_session);

    //Parse tempFile

    string line;
    ifstream myfile ("tempFile.txt");
    if (myfile.is_open())
    {
        int commentFound=0;
        int lineIndex=0;
        while ( myfile.good() )
        {
            //If a line has comment then we have a small problem
            getline (myfile,line);
            if (lineIndex>=2)
            {
                istringstream iss(line);
                if (commentFound==1) //second field is the ip address
                {
                    string t1,t2;
                    iss >> t1;iss >> t2;
                    string IP=t2;
                    retVal.push_back(IP);
                    commentFound=0;
                }
                else //third field is the ip address if no comment present
                {
                    string t1,t2,t3;
                    iss >> t1;iss >> t2;iss >> t3;
                    if (t2.compare(";;;")==0)
                    {
                        commentFound=1;
                    }
                    else
                    {
                        string IP=t3;
                        retVal.push_back(IP);
                    }

                }
            }
            else
            {
                lineIndex++;
            }
        }
        myfile.close();
    }

    return retVal;
}

int addIPToFirewall(string IP)
{
    int retVal=0;

    ssh_session my_ssh_session;
    int rc;
    my_ssh_session = ssh_new();
    if (my_ssh_session == NULL)
        return -1;
    ssh_options_set(my_ssh_session, SSH_OPTIONS_HOST, conf->getFirewallIP().c_str());
    ssh_options_set(my_ssh_session, SSH_OPTIONS_USER, conf->getFirewallUser().c_str());
    rc = ssh_connect(my_ssh_session);
    if (rc != SSH_OK)
    {
        string msg="SSH Connection error";
        log(1,msg,0);
        return -2;
    }
    rc = ssh_userauth_password(my_ssh_session, NULL, conf->getFirewallPass().c_str());
    if (rc != SSH_AUTH_SUCCESS)
    {
        string msg="SSH authentication error";
        log(1,msg,0);
        return -3;
    }
    ssh_channel channel=ssh_channel_new(my_ssh_session);
    rc = ssh_channel_open_session(channel);
    if (rc != SSH_OK)
    {
        string msg="SSH channel error";
        log(1,msg,0);
        ssh_channel_free(channel);
        return -4;
    }

    string cmd="/ip firewall address-list add list=" + conf->getFirewallAddressList() + " address=" + IP;
    rc = ssh_channel_request_exec(channel, cmd.c_str());
    if (rc != SSH_OK)
    {
        string msg="SSH command execution error (" + cmd + ")";
        log(1,msg,0);
        ssh_channel_close(channel);
        ssh_channel_free(channel);
        return -5;
    }
    else
    {
        string msg="SSH executed : " + cmd;
        log(3,msg,1);
    }
    ssh_channel_send_eof(channel);
    ssh_channel_close(channel);
    ssh_channel_free(channel);
    ssh_disconnect(my_ssh_session);
    ssh_free(my_ssh_session);
    retVal=1;
    return retVal;
}

int removeIPFromFirewall(string IP)
{
    ssh_session my_ssh_session;
    int rc;

    my_ssh_session = ssh_new();
    if (my_ssh_session == NULL)
        return -1;

    ssh_options_set(my_ssh_session, SSH_OPTIONS_HOST, conf->getFirewallIP().c_str());
    ssh_options_set(my_ssh_session, SSH_OPTIONS_USER, conf->getFirewallUser().c_str());

    rc = ssh_connect(my_ssh_session);
    if (rc != SSH_OK)
    {
        string msg="SSH Connection error";
        log(1,msg,0);
        return -2;
    }

    rc = ssh_userauth_password(my_ssh_session, NULL, conf->getFirewallPass().c_str());
    if (rc != SSH_AUTH_SUCCESS)
    {
        string msg="SSH authentication error";
        log(1,msg,0);
        return -3;
    }

    ssh_channel channel=ssh_channel_new(my_ssh_session);
    rc = ssh_channel_open_session(channel);
    if (rc != SSH_OK)
    {
        string msg="SSH channel error";
        log(1,msg,0);
        ssh_channel_free(channel);
        return -4;
    }
    string cmd="/ip firewall address-list remove [/ip firewall address-list find list=" + conf->getFirewallAddressList() + " address=" + IP + "]";

    rc = ssh_channel_request_exec(channel, cmd.c_str());

    if (rc != SSH_OK)
    {
        string msg="SSH command execution error (" + cmd + ")";
        log(1,msg,0);
        ssh_channel_close(channel);
        ssh_channel_free(channel);
        return -5;
    }
    else
    {
        string msg="SSH executed : " + cmd;
        log(3,msg,1);
    }

    ssh_channel_send_eof(channel);
    ssh_channel_close(channel);
    ssh_channel_free(channel);


    ssh_disconnect(my_ssh_session);
    ssh_free(my_ssh_session);

}

void *leaseRemover(void *ptr)
{
    while(1)
    {
        attackleases::attackLeasesStructure tmp=al->getLease(0);
        struct timeb tp;
        ftime(&tp);

        while ((tmp.IP.compare("")!=0) && (tp.time-tmp.tm>=conf->getAttackLeaseTimeout()))
        {
            al->removeLease(0);
            removeIPFromFirewall(tmp.IP);
            string msg="Removed Lease for : " + tmp.IP;
            log(2,msg,1);
            tmp=al->getLease(0);
        }

        if ((tmp.tm>0) && (conf->getAttackLeaseTimeout() + tmp.tm - tp.time>=0))
        {
            string msg="leaseRemover sleeping for " + itos(conf->getAttackLeaseTimeout() + tmp.tm - tp.time) + " seconds";
            log(4,msg,1);
            sleep(conf->getAttackLeaseTimeout() + tmp.tm - tp.time);
        }
        else
        {
            string msg="leaseRemover sleeping for " + itos(conf->getAttackLeaseTimeout()) + " seconds";
            log(4,msg,1);
            sleep(conf->getAttackLeaseTimeout());
        }
    }
}

void *attackBlocker(void *ptr)
{

    while(1)
    {       
        pthread_mutex_lock(&mutex2);
        while (newAttackAdded==0)
        {
            pthread_cond_wait(&newAttackCond, &mutex2);
        }
        newAttackAdded=0;
        pthread_mutex_unlock(&mutex2);

        struct timeb tp;
        ftime(&tp);
        int i=0;
        int atz=attacker.size();
        for (i=0;i<atz;i++)
        {
            string ts=attacker[i];
            if (ts.empty()==false)
            {
                if (al->leaseExist(ts)==0)
                {
                    ftime(&tp);
                    al->addLease(ts,tp.time);
                    attacker.erase(attacker.begin()+i);
                    i--;
                    addIPToFirewall(ts);
                    string msg="Lease added for : " + ts;
                    log(2,msg,1);
                }
            }
        }
    }

}

void removeIPFromConnections(string IP)
{
    int i=0;
    connection::connectionStructure cnx=c->getConnection(i++);


    while (cnx.tm>0)
    {
        if (cnx.srcIP.compare(IP)==0)
        {
            c->removeConnection(i);
            i--;
        }
        cnx=c->getConnection(i++);
    }
}

void *attackChecker(void *ptr)
{
    while(1)
    {
        pthread_mutex_lock(&mutex1);
        while (newConnectionAdded==0)
        {
            pthread_cond_wait(&newConnectionCond, &mutex1);
        }
        newConnectionAdded=0;
        pthread_mutex_unlock(&mutex1);
        vector<string> parsedIPs;
        int i=0;

        connection::connectionStructure tmp=c->getConnection(i);
        while (tmp.srcIP.compare("")!=0)
        {
            int j=0;
            int found=0;

            for (j=0;j<parsedIPs.size();j++)
            {
                string ts=parsedIPs[j];
                if (ts.compare(tmp.srcIP)==0)
                {
                    found=1;
                    break;
                }
            }

            if (found==0)
            {
               parsedIPs.push_back(tmp.srcIP);
               int k=i+1;
               connection::connectionStructure tmp1=c->getConnection(k);
               int sc=0;
               while (tmp1.srcIP.compare("")!=0)
               {
                   if (tmp1.srcIP.compare(tmp.srcIP)==0) //Compare Source IPs
                   {
                       if (tmp1.destIP.compare(tmp.destIP)==0) //Compare Destination IPs
                       {
                           if (tmp1.destPort==tmp.destPort) //Compare Destination Ports
                           {
                               //Probably second packet
                               sc=sc+1;
                           }
                           else
                           {
                               //Probably port scanner
                               sc=sc+4;
                           }
                       }
                       else //Same source different destination
                       {
                           if (tmp1.destPort==tmp.destPort)
                           {
                               //Probably ip scanner
                               sc=sc+4;
                           }
                           else //Unknown attack
                           {
                               sc=sc+3;
                           }
                       }

                   }

                   k++;
                   tmp1=c->getConnection(k);
               }

               if (sc>=conf->getScannerThreshold())
               {
                   int rfound=0;
                   for (int r=0;r<attacker.size();r++)
                   {
                       string tas=attacker[r];
                       if (tas.compare(tmp.srcIP)==0)
                       {
                           rfound=1;
                           break;
                       }
                   }
                   if (rfound==0)
                   {
                       attacker.push_back(tmp.srcIP);
                       removeIPFromConnections(tmp.srcIP);
                       string msg="Attacker detected : " + tmp.srcIP + "  SCORE=" + itos(sc);
                       log(3,msg,1);
                       newAttackAdded=1;
                       pthread_cond_signal(&newAttackCond);
                   }
               }
            }

            i++;
            tmp=c->getConnection(i);
        }
    }

}

void *oldConnectionChecker(void *ptr)
{
    while(1)
    {
        connection::connectionStructure cnx=c->getConnection(0);
        struct timeb tp;
        ftime(&tp);
        while ((cnx.tm>0) && (tp.time-cnx.tm>=connectionTimeout))
        {
            string msg="oldConnectionChecker removing connection (" + cnx.srcIP + ":" + itos(cnx.srcPort) + "->" + cnx.destIP + ":" + itos(cnx.destPort) + ")";
            log(4,msg,1);
            c->removeConnection(0);
            cnx=c->getConnection(0);
        }

        if ((cnx.tm>0) && (connectionTimeout + cnx.tm - tp.time>=0))
        {
            string msg="oldConnectionChecker sleeping for " + itos(connectionTimeout + cnx.tm - tp.time) + " seconds";
            log(4,msg,1);
            sleep(connectionTimeout + cnx.tm - tp.time);
        }
        else
        {
            string msg="oldConnectionChecker sleeping for " + itos(connectionTimeout) + " seconds";
            log(4,msg,1);
            sleep(connectionTimeout);
        }
    }
}

void signal_callback_handler(int signum)
{
    //Delete pid file
    remove(conf->getPidFile().c_str());
    string msg="stopped successfully";
    log(1,msg,1);
    exit(signum);
}

int main(int argc, char *argv[])
{
    conf=new configfilereader(configFile);
    string confCheck=conf->checkOptions();

    if (confCheck.compare("OK")!=0)
    {
        cout << "Error in configuration file : " << confCheck << endl;
        string msg="Error in configuration file : " + confCheck;
        log(1,msg,0);
        exit(-1);
    }

    if (daemon(0, 0) >= 0)
    {
        string msg="daemonize OK";
        log(3,msg,1);
    }
    else
    {
        string msg="Could not daemonize";
        log(1,msg,0);
        exit(-3);
    }

    setsid();
    ofstream pf;
    pf.open (conf->getPidFile().c_str(), ios::out);
    if (pf.is_open())
    {
        pf << getpid();
        pf.close();
    }
    else
    {
        string msg="Could not open pid file " + conf->getPidFile();
        log(1,msg,0);
    }

    c=new connection();
    al=new attackleases(conf->getLeasesFile());

    int  iret1, iret2, iret3, iret4;
    iret1 = pthread_create( &attackCheckerThread, NULL, attackChecker, NULL);
    iret2 = pthread_create( &attackBlockerThread, NULL, attackBlocker, NULL);
    iret3 = pthread_create( &oldConnectionCheckerThread, NULL, oldConnectionChecker, NULL);
    iret4 = pthread_create( &leaseRemoverThread, NULL, leaseRemover, NULL);

    struct sockaddr_in my_addr, cli_addr;
    int sockfd;
    socklen_t slen=sizeof(cli_addr);
    char buf[BUFLEN];

    if ((sockfd = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP))==-1)
    {
        string msg="Could not open socket";
        log(1,msg,0);
        exit(-1);
    }
    else
    {
        string msg="Socket opened successfully";
        log(1,msg,1);
    }

    bzero(&my_addr, sizeof(my_addr));
    my_addr.sin_family = AF_INET;
    my_addr.sin_port = htons(conf->getListenPort());
    my_addr.sin_addr.s_addr = htonl(INADDR_ANY);

    if (bind(sockfd, (struct sockaddr* ) &my_addr, sizeof(my_addr))==-1)
    {
        string msg="Could not bind to socket";
        log(0,msg,1);
        exit(-2);
    }
    else
    {
        string msg="bind successful";
        log(1,msg,1);
    }

    string msg="started successfully (pid=" + itos(getpid()) + ")";
    log(1,msg,1);

    signal(SIGINT, signal_callback_handler);
    signal(SIGHUP, signal_callback_handler);

    while(1)
    {
        int br=recvfrom(sockfd, buf, BUFLEN, 0, (struct sockaddr*)&cli_addr, &slen);
        string *ca= new string (inet_ntoa(cli_addr.sin_addr));

        if (ca->compare(conf->getFirewallIP())==0)
        {
            if (br==-1)  cout << "recvfrom()";
            if (br<BUFLEN) buf[br]='\0';
            string sbuf=string(buf);
            string srest,protocol,srcIP,srcPort,destIP,destPort;
            size_t found;
            found=sbuf.find("proto");
            srest=sbuf.substr(found+6);
            protocol=srest.substr(0,3);
            if ( (protocol.compare("UDP")==0) || (protocol.compare("TCP")==0) )
            {
                found=srest.find(", ");
                srest=srest.substr(found+2);
                found=srest.find(":");
                srcIP=srest.substr(0,found);
                srest=srest.substr(found+1);
                found=srest.find("-");
                srcPort=srest.substr(0,found);
                srest=srest.substr(found+2);
                found=srest.find(":");
                destIP=srest.substr(0,found);
                srest=srest.substr(found+1);
                found=srest.find(",");
                destPort=srest.substr(0,found);
                struct timeb tp;
                ftime(&tp);
                c->addConnection(protocol,srcIP,srcPort,destIP,destPort,tp.time);
                string msg="New connection added : " + srcIP + ":" + srcPort + "->" + destIP + ":" + destPort;
                log(4,msg,1);
                newConnectionAdded=1;
                pthread_cond_signal(&newConnectionCond);
            }
            else
            {
                string msg="Ignoring unknown protocol";
                log(4,msg,1);
            }
        }
        else
        {
            string msg="Ignoring remote server with ip : " + *ca;
            log(3,msg,1);
        }
    }

    close(sockfd);

    return 0;
}
